
/*
 *
 * Copyright 2003 Blur Studio Inc.
 *
 * This file is part of libstone.
 *
 * libstone is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * libstone is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with libstone; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */

/*
 * $Id: ffimagesequenceprovider.cpp 5411 2007-12-18 01:03:08Z brobison $
 */

#ifdef USE_FFMPEG

#define __STDC_CONSTANT_MACROS
#define __STDC_LIMIT_MACROS

#include <stdint.h>
#include <inttypes.h>
#include <stdio.h>
#include "ffimagesequenceprovider.h"

extern "C" {
#include "libavformat/avformat.h"
#include "libavcodec/avcodec.h"
#include "libswscale/swscale.h"
#include "libavutil/md5.h"
#include "libavutil/log.h"
}

#include "blurqt.h"

class FFLoader
{
public:
	FFLoader();
	
	bool open( const QString & fileName );
	bool seek( int frameNumber );
	QImage readImage( int frameNumber );

protected:
	bool convertToRGB( AVFrame * frame, QImage * );
	QImage convertToQImage( AVFrame * frame );

	int64_t frameNumberToTimestamp( int frameNumber );

	AVFrame * readNextFrame();

	int mFrameStart, mFrameEnd;
	double mFrameDuration;
	FFImageSequenceProvider * mProvider;
	AVFormatContext * mFormatContext;
	AVCodecContext * mCodecContext;
	AVCodec * mCodec;
	AVDictionary * mCodecOpts;
	AVFrame * mFrame;
	int mStreamIndex;
	AVStream * mStream;
	bool mIsOpen;
	AVPacket * mPacket;
	int mPacketDataPos;
	int mLastFrameRead;
	QImage mLastImage;
	SwsContext * mImgConvertCtx;
	
	QMutex mLoadedMutex;
	QList<QPair<int,QImage> > mLoaded;
};

class FFLoaderThread : public QThread
{
public
	FFLoaderThread( QObject * parent, FFLoader * loader );
	
	FFLoader * mLoader;
};

class FFImageSequenceProvider::Private
{
public:
	Private( FFImageSequenceProvider * provider );
	~Private();

	bool open( const QString & fileName );
	bool seek( int frameNumber );
	QImage readImage( int frameNumber );

protected:

	FFLoaderThread * mLoaderThread;
};

FFImageSequenceProvider::Private::Private( FFImageSequenceProvider * provider )
: mProvider( provider )
, mFormatContext( 0 )
, mCodecContext( 0 )
, mCodec( 0 )
, mCodecOpts( 0 )
, mFrame( 0 )
, mStreamIndex( -1 )
, mStream( 0 )
, mIsOpen( false )
, mPacket( 0 )
, mPacketDataPos( 0 )
, mLastFrameRead( -1 )
, mImgConvertCtx( 0 )
{}

FFImageSequenceProvider::Private::~Private()
{
	if( mPacket )
		av_free_packet( mPacket );
	if( mCodecContext ) {
		avcodec_close( mCodecContext );
		mCodecContext = 0;
	}
	if( mFormatContext ) {
		avformat_close_input( &mFormatContext );
		mFormatContext = 0;
	}
	if( mImgConvertCtx ) {
		sws_freeContext( mImgConvertCtx );
		mImgConvertCtx = 0;
	}
	if( mFrame ) {
		av_free(mFrame);
	}
}

bool FFImageSequenceProvider::Private::open( const QString & fileName )
{
	QByteArray fileNameLatin1 = fileName.toLatin1();
	
	av_log_set_level(AV_LOG_DEBUG);
	//mFormatContext = new AVFormatContext();
	
	if( avformat_open_input( &mFormatContext, fileNameLatin1.constData(), NULL, NULL ) ) {
		LOG_1( "avformat_open_input failed to open file: " + fileName );
		return false;
	}

	if( avformat_find_stream_info( mFormatContext, NULL ) < 0 ) {
		LOG_1( "av_find_stream_info failed to find stream info for file: " + fileName );
		return false;
	}

#ifdef AV_DEBUG
	dump_format( mFormatContext, 0, fileNameLatin1.constData(), false );
#endif

	mStreamIndex = av_find_best_stream( mFormatContext, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0 );
	mStream = mFormatContext->streams[mStreamIndex];
	mCodecContext = mStream->codec;

	mCodec = avcodec_find_decoder( mCodecContext->codec_id );
	
	if(mCodec->capabilities & CODEC_CAP_DR1)
		mCodecContext->flags |= CODEC_FLAG_EMU_EDGE;
	
	if( !mCodec ) {
		LOG_5( "Unable to find decoder for stream" );
		return false;
	}
	
	// Inform the codec that we can handle truncated bitstreams -- i.e.,
	// bitstreams where frame boundaries can fall in the middle of packets
	if( mCodec->capabilities & CODEC_CAP_TRUNCATED )
		mCodecContext->flags |= CODEC_FLAG_TRUNCATED;
	
	// Open codec
	if( avcodec_open2(mCodecContext, mCodec, NULL) < 0 ) {
		LOG_5( "Unable to open codec" );
		return false;
	}

	// Hack to correct wrong frame rates that seem to be generated by some 
	// codecs
	if( mStream->r_frame_rate.num > 1000 && mStream->r_frame_rate.den == 1 )
		mStream->r_frame_rate.den=1000;

	mFrameDuration = 1000.0 * mStream->r_frame_rate.den / mStream->r_frame_rate.num;
	LOG_5( "Frame Duration: " + QString::number( mFrameDuration ) + "ms" );

	mFrameStart = (mStream->start_time * mStream->time_base.num * mStream->r_frame_rate.num) / (mStream->time_base.den * mStream->r_frame_rate.den);
	mFrameEnd = ((mStream->start_time + mStream->duration) * mStream->time_base.num * mStream->r_frame_rate.num) / (mStream->time_base.den * mStream->r_frame_rate.den);
	mLastFrameRead = mFrameStart - 1;

	mIsOpen = true;
	return true;
}

int64_t FFImageSequenceProvider::Private::frameNumberToTimestamp( int frameNumber )
{
	// seconds = frameNumber / frameRate
	// timestamp = seconds / time_base
	// timestamp = (frameNumber / frameRate) / time_base
	// timestamp = frameNumber / (time_base * frameRate)
	return (int64_t(frameNumber) * mStream->time_base.den * mStream->r_frame_rate.den) / (int64_t(mStream->time_base.num) * mStream->r_frame_rate.num);
}

bool FFImageSequenceProvider::Private::seek( int frameNumber )
{
	mLastFrameRead = frameNumber - 1;
	int64_t timestamp = frameNumberToTimestamp( frameNumber );
	LOG_5( "Seeking frame " + QString::number( frameNumber ) );
	if( mPacket ) {
		av_free_packet( mPacket );
		delete mPacket;
		mPacket = 0;
		mPacketDataPos = 0;
	}
	if( avformat_seek_file( mFormatContext, mStreamIndex, INT64_MIN, timestamp, INT64_MAX, 0 /*AVSEEK_FLAG_FRAME*/ ) < 0 ) {
		LOG_1( "Failed to seek to frame" + QString::number(frameNumber) );
		return false;
	}
	
	avcodec_flush_buffers(mCodecContext);
	
	return true;
}

const char * packetMd5( AVPacket * pkt )
{
	static char out[35];
	uint8_t md5sum[16];
	av_md5_sum(md5sum, (const uint8_t*)(pkt->data), pkt->size);
	out[0] = '0';
	out[1] = 'x';
	for( int i=0; i<16; ++i )
		sprintf( &out[i+2], "%X%X", md5sum[i] & 0xf, md5sum[i] >> 4 );
	return out;
}

AVFrame * FFImageSequenceProvider::Private::readNextFrame()
{
	if( !mFrame )
		mFrame = avcodec_alloc_frame();

	int frameFinished = 0, err = 0;

	mLastFrameRead++;

	if( !mPacket ) {
		mPacket = new AVPacket;
		av_init_packet(mPacket);
		//mPacket->size = 0;
	}
	
	while( frameFinished == 0 ) {
		
		if( mPacketDataPos >= mPacket->size ) {
			av_free_packet(mPacket);
			mPacketDataPos = 0;
			while (true) {
				if (av_read_frame(mFormatContext, mPacket) < 0) {
					LOG_1( "av_read_frame returned error code: " + QString::number(err) );
					return 0;
				}
				if( mPacket->stream_index == mStreamIndex ) break;
				av_free_packet(mPacket);
			}
		}
		
		// preparing the packet that we will send to libavcodec
		AVPacket packetToSend;
		av_init_packet(&packetToSend);
		packetToSend.data = mPacket->data + mPacketDataPos;
		packetToSend.size = mPacket->size - mPacketDataPos;

		// sending data to libavcodec
		int processedLength = avcodec_decode_video2( mCodecContext, mFrame, &frameFinished, &packetToSend);
		
		if (processedLength < 0) {
			av_free_packet(mPacket);
			return 0;
		}
		
		LOG_5( QString("Processed %1 bytes of %2 in packet").arg(processedLength).arg(mPacket->size) );
		mPacketDataPos += processedLength;
	}
	
	return mFrame;
}

bool FFImageSequenceProvider::Private::convertToRGB( AVFrame * frame, QImage * image )
{
	mImgConvertCtx = sws_getCachedContext( mImgConvertCtx, mCodecContext->width, mCodecContext->height, mCodecContext->pix_fmt,
		mCodecContext->width, mCodecContext->height, PIX_FMT_BGRA, SWS_FAST_BILINEAR, NULL, NULL, NULL);

	if (mImgConvertCtx == NULL)
	{
		LOG_1( "Failed to get SwsContext for image format conversion.");
		return false;
	}
	
	uint8_t * imgPlanes [] = { image->bits(), image->bits(), image->bits(), image->bits() };
	int strides [] = { image->bytesPerLine(), image->bytesPerLine(), image->bytesPerLine(), image->bytesPerLine() };
	return sws_scale( mImgConvertCtx, frame->data, frame->linesize, 0, mCodecContext->height, imgPlanes, strides) == image->height();
}

QImage FFImageSequenceProvider::Private::convertToQImage( AVFrame * frame )
{
	QImage img( mCodecContext->width, mCodecContext->height, QImage::Format_ARGB32 );

	if( mCodecContext->pix_fmt != PIX_FMT_ARGB ) {
		if( !convertToRGB( frame, &img ) ) {
			LOG_1( "convertToRGB failed" );
		}
	} else
		memcpy( img.bits(), frame->data[0], img.width() * img.height() * 4 );

/*	for( int i = 0; i < img.width() * img.height(); i++ ) {
		QRgb c =  img.pixel( i % img.width(), i / img.width() );
		if( qRed(c) + qGreen(c) + qBlue(c) > 0 ) {
			//LOG_5( "Image not all black" );
			break;
		}
	} */

	return img;
}

QImage FFImageSequenceProvider::Private::readImage( int frameNumber )
{
	if( mLastFrameRead != frameNumber - 1 ) {

		// Return the cached image
		if( mLastFrameRead == frameNumber && !mLastImage.isNull() )
			return mLastImage;

		if( !seek( frameNumber ) ) {
			LOG_5( "Seek failed for frame: " + QString::number( frameNumber ) );
			return QImage();
		}
	}

	AVFrame * frame = readNextFrame();
	if( !frame ) return QImage();
	QImage img = convertToQImage( frame );
	mLastImage = img;
	return img;
}


FFImageSequenceProvider::FFImageSequenceProvider( const QString & path, QObject * parent )
: ImageSequenceProvider( parent )
, d( new Private( this ) )
, mFilePath( path )
{
	d->open( path );
}

FFImageSequenceProvider::~FFImageSequenceProvider()
{
	delete d;
}

bool FFImageSequenceProvider::isOpen()
{
	return d->mIsOpen;
}

QString FFImageSequenceProvider::path()
{
	return mFilePath;
}

int FFImageSequenceProvider::frameStart()
{
	return d->mFrameStart;
}

int FFImageSequenceProvider::frameEnd()
{
	return d->mFrameEnd;
}

ImageSequenceProvider::ImageStatus FFImageSequenceProvider::status( int )
{
	// Assume all images in a video are available
	return ImageAvailable;
}

QImage FFImageSequenceProvider::image( int frameNumber )
{
	return d->readImage( frameNumber );
}


FFImageSequenceProviderPlugin::FFImageSequenceProviderPlugin()
{
	av_register_all();
}

FFImageSequenceProviderPlugin::~FFImageSequenceProviderPlugin(){}

QStringList FFImageSequenceProviderPlugin::fileExtensions()
{
	QStringList ret;
	AVInputFormat * format = 0;
	while( format = av_iformat_next(format) ) {
		//fprintf( stderr, "avformat AVInputFormat: Name: %s Long Name %s Extensions %s\n", format->name, format->long_name, format->extensions );
		if( format->extensions )
			ret += QString::fromLatin1( format->extensions ).split(',');
//		ret += QString::fromLatin1( 
	}
	ret << "mpg" << "mpeg" << "avi";
	ret.sort();
	//LOG_5( "Supported Extensions: " + ret.join(",") );
	return ret;
}

// Could use a timer to delete this if it's not used
static FFImageSequenceProvider * nextProvider = 0;

bool FFImageSequenceProviderPlugin::supportsFormat( const QString & fileName )
{
	if( nextProvider ) delete nextProvider;
	nextProvider = new FFImageSequenceProvider( fileName );
	return nextProvider->isOpen();
}

ImageSequenceProvider * FFImageSequenceProviderPlugin::createProvider( const QString & fileName )
{
	if( nextProvider && nextProvider->path() == fileName ) {
		ImageSequenceProvider * ret = nextProvider;
		nextProvider = 0;
		return ret;
	}
	return new FFImageSequenceProvider( fileName );
}

#endif // USE_FFMPEG

void registerFFImageSequenceProviderPlugin()
{
#ifdef USE_FFMPEG
	FFImageSequenceProviderPlugin * plugin = new FFImageSequenceProviderPlugin();
	plugin->fileExtensions();
	ImageSequenceProviderFactory::registerPlugin( plugin );
#endif // USE_FFMPEG
}

